import os
import traceback
import base64
from datetime import datetime
from typing import Union, Sequence, Optional, Set
from mispconn.utils import logger, CFG, parse_excluding
from tools import Tools


class AttachmentDownloader:
    """Download attachments from MISP"""
    def __init__(self, attribute_type: str, path: str, tags: str, num: int,
                 file_paths_for_excluding: str):
        """
        :param attribute_type: MISP attribute type
        :param path: path for saving downloads
        :param tags: misp tags on both events and attributes
        :param num: number of files
        :param file_paths_for_excluding: ignores attribute ids in the file
        """
        self.attribute_type = attribute_type
        self.path = path
        self.num = num
        self.file_path_for_future_excluding: Optional[str] = None
        self.excluded_attrs: Set[str] = set()
        self.init_excluding(file_paths_for_excluding)
        assert self.attribute_type and self.path
        # for building build_complex_query
        self.tags = {'or': [], 'and': [], 'not': []}
        if tags:
            if isinstance(tags, list):
                tags = tags
            else:
                tags = [tags]
        else:
            tags = []
        for tag in tags:
            if tag.startswith('or:'):
                self.tags['or'].append(tag[3:])
            elif tag.startswith('and:'):
                self.tags['and'].append(tag[4:])
            elif tag.startswith('not:'):
                self.tags['not'].append(tag[4:])

    def execute(self):
        misp = Tools.init_misp(CFG['misp']['url'], CFG['misp']['key'])
        try:
            res = misp.search(
                controller='attributes',
                type_attribute=self.attribute_type,
                # tags=self.tags,
                tags=misp.build_complex_query(and_parameters=self.tags['and'],
                                              not_parameters=self.tags['not'],
                                              or_parameters=self.tags['or']),
                # tags=misp.build_complex_query(and_parameters=self.tags, not_parameters=['secheart:reputation="malware"']),
                include_correlations=False,
                with_attachments=False,
                metadata=True,
                limit=self.num)
            attrs = res['Attribute']
            logger.info(f'found {len(attrs)} attributes in misp.')
            for attr in attrs:
                if attr['id'] in self.excluded_attrs:
                    logger.debug(f'{attr["id"]} was already downloaded, '
                                 f'ignores.')
                    continue
                name = attr['value']
                path = os.path.join(self.path, name)
                data = misp.get_attribute(attr['id'])['Attribute']['data']
                with open(path, 'wb') as f:
                    f.write(base64.b64decode(data))
                self.add_to_file_for_excluding(attr['id'])
                logger.debug(f'downloaded {path} successfully.')
            logger.info('all attributes are downloaded, done')
        except Exception as e:
            logger.critical(e)
            logger.debug(traceback.print_exc())

    def init_excluding(self, file_paths_for_excluding:
    Union[Sequence[str], str, None]):
        # Create a new file to record downloaded attribute ids for this run,
        # which can be used in the future runs.
        fname = 'aicti_pf_excluded_attrs_in' + \
                self.path.replace('/', '_') + '_' + \
                datetime.now().strftime('%Y%m%d%H%M%S')
        self.file_path_for_future_excluding = os.path.join('/tmp', fname)
        logger.info(f'processed attribute ids will be recorded in '
                    f'{self.file_path_for_future_excluding} for future '
                    f'exclusion.')
        self.excluded_attrs = parse_excluding(file_paths_for_excluding)

    def add_to_file_for_excluding(self, file_path: str):
        with open(self.file_path_for_future_excluding, 'a') as f:
            f.write(f'{file_path}\n')
